from utils.basic_function import *
from utils.database_table import FusedRoadTable, RSNodeTable


class StakeTool:
    def __init__(self, road_table, node_table, east_west_angle):
        self.road_table = road_table
        self.node_table = node_table
        self.east_west_angle = np.abs(east_west_angle)
        self.road_table.add_column(self.road_table.stake_rule, set_null=False)

    @staticmethod
    def _stake_interpolate_simple(linestring: LineString, pt: Point, start_stake, end_stake):
        """
        不考虑桩号编制原则，从linestring的起点开始计算。
        """
        actual_length = linestring.length
        line_coord_list = list(linestring.coords)
        before_length = 0
        proj_result = project_pt_to_polyline(pt.coords[0], line_coord_list)
        for i in range(proj_result["pre_node"]):
            before_length += cal_dist(line_coord_list[i], line_coord_list[i + 1])
        before_length += cal_dist(line_coord_list[proj_result["pre_node"]], proj_result["proj_pt"])
        stake = start_stake + before_length / actual_length * (end_stake - start_stake)
        return stake

    def label_stake_direction(self, road_pk):
        """
        考虑桩号编制原则，判断道路起点与桩号的对应关系。
        """
        gdf = geopd.read_postgis(f"select * from {self.road_table.name} where {self.road_table.pk} = {road_pk}", conn)
        road_code = gdf.loc[0, FusedRoadTable.road_code]
        road_geom = gdf.loc[0, FusedRoadTable.geom]
        snodeid = gdf.loc[0, FusedRoadTable.snodeid]
        enodeid = gdf.loc[0, FusedRoadTable.enodeid]
        start_stake = gdf.loc[0, FusedRoadTable.start_stake]
        end_stake = gdf.loc[0, FusedRoadTable.end_stake]

        # 新增路网
        if road_code[0] == 'N':
            start_idx, stake_rule = self.get_new_road_start_stake_idx(road_pk, road_geom, snodeid, enodeid)
            if start_idx == -1:  # 需要转换方向
                road_geom = reverse_linestring(road_geom)
                snodeid, enodeid = enodeid, snodeid
            self.road_table.update_attr(self.road_table.pk, road_pk,
                                        {self.road_table.stake_rule: stake_rule, self.road_table.geom: road_geom,
                                         self.road_table.snodeid: snodeid, self.road_table.enodeid: enodeid})
            return

        # 上报路网
        else:
            reverse_flag = False
            # 从属性表中读取起点和终点桩号
            # 根据相邻的同路线路段判断snode、enode桩号
            snode_stake = self.get_old_node_stake_by_neighbor(road_code, snodeid)
            enode_stake = self.get_old_node_stake_by_neighbor(road_code, enodeid)

            if snode_stake is None and enode_stake is None:  # 该路线只有一条路段，从北向南
                stake_rule = "NS"
                direction = self.compute_road_direction(road_geom.coords[0], road_geom.coords[-1])
                if 0 <= direction < 180:  # 由南向北，需要翻转
                    reverse_flag = True
            else:  # 该路线有多条路段，根据与其他路段的关系判断起点
                stake_rule = "multi_segment"
                if snode_stake is not None:
                    if snode_stake == end_stake:  # 几何起点对应终点桩号，需要翻转
                        reverse_flag = True
                elif enode_stake is not None:
                    if enode_stake == start_stake:
                        reverse_flag = True
            if reverse_flag:  # 转换方向
                road_geom = reverse_linestring(road_geom)
                snodeid, enodeid = enodeid, snodeid
            self.road_table.update_attr(self.road_table.pk, road_pk,
                                        {self.road_table.stake_rule: stake_rule, self.road_table.geom: road_geom,
                                         self.road_table.snodeid: snodeid, self.road_table.enodeid: enodeid})

    def stake_interpolate(self, ldbm, pt):
        gdf = geopd.read_postgis(f"select * from {self.road_table.name} where {self.road_table.segment_code} = '{ldbm}'", conn)
        linestring = gdf.loc[0, GEOM_COL]
        start_stake = gdf.loc[0, self.road_table.start_stake]
        end_stake = gdf.loc[0, self.road_table.end_stake]
        return self._stake_interpolate_simple(linestring, pt, start_stake, end_stake)

    def get_old_node_stake_by_neighbor(self, road_code, node_id):
        """
        根据同一路线的多条路段的排列顺序判断结点桩号。

        Args:
            road_code: 路线编码
            node_id: 结点编号

        Returns:
            若node_id所指结点连接了同一路线的不同路段，则根据相邻路段的公共桩号判断node_id的桩号并返回；
            否则返回None，即无法根据相邻情况判断桩号
        """
        df = pd.read_sql(f"""select * from {self.road_table.name} where {FusedRoadTable.road_code} = '{road_code}' and
            ({FusedRoadTable.snodeid}={node_id} or {FusedRoadTable.enodeid}={node_id})""", conn)
        if df.shape[0] < 2:  # 该结点没有同路线的相邻路段，无法判断该结点的桩号
            return None
        # 四个桩号中，排在中间的两个相同的桩号即为该结点的桩号
        stake_list = [df.loc[0, FusedRoadTable.start_stake], df.loc[0, FusedRoadTable.end_stake],
                      df.loc[1, FusedRoadTable.start_stake], df.loc[1, FusedRoadTable.end_stake]]
        # 桩号排序，验证中间两个桩号相等
        stake_list.sort()
        if stake_list[1] != stake_list[2]:
            raise Exception(f"{self.road_table.name}中结点{self.node_table.pk}={node_id}连接的路段的桩号不一致")
        return stake_list[1]

    def get_new_road_start_stake_idx(self, pk_val, geom, snodeid, enodeid):
        """
        判断新增道路起点桩号对应几何线形的起点还是终点。
        判断原则：
        1.当单线段路连接高等级路和低等级路时，路线指向为高等级向低等级路
        2.当单线段路连接同等级公路时，路线指向为由北向南
        3.当单线段路只有一侧连接一条路，另一侧没有连接时，路线指向为由北向南、由东向西（首要遵守由北向南）

        Returns:
            start_idx: 几何线形的起点对应起点桩号时返回0，否则返回-1；
            stake_rule: str
        """
        # 判断端点情况
        snode_class = self.get_node_highest_class(pk_val, snodeid)
        enode_class = self.get_node_highest_class(pk_val, enodeid)

        # 两侧都连接道路
        if snode_class is not None and enode_class is not None:
            node_rank = compare_road_rank(snode_class, enode_class)
            # case 1：两侧道路等级不同，从高级到低级
            if node_rank == 1:  # snode的等级比enode高，符合从高等级指向低等级
                return 0, "two_side_HL"
            elif node_rank == -1:  # snode等级低于enode，需要翻转
                return -1, "two_side_HL"
            # case 2：两侧同等级，由北向南
            else:
                direction = self.compute_road_direction(geom.coords[0], geom.coords[-1])
                if -180 <= direction < 0:  # 符合由北向南
                    return 0, "two_side_NS"
                else:  # 由南向北，需要翻转
                    return -1, "two_side_NS"
        # case 3：仅一侧连接其他道路或完全孤立，由北向南，由东向西
        else:
            direction = self.compute_road_direction(geom.coords[0], geom.coords[-1])
            if -180 + self.east_west_angle <= direction <= -self.east_west_angle:  # 符合由北向南
                return 0, "dangle_NS"
            elif self.east_west_angle <= direction <= 180 - self.east_west_angle:  # 由南向北，需要翻转
                return -1, "dangle_NS"
            elif np.abs(direction) < self.east_west_angle:  # 由西向东，需要翻转
                return -1, "dangle_EW"
            else:  # 符合由东向西
                return 0, "dangle_EW"

    @staticmethod
    def compute_road_direction(st, ed):
        """
        计算向量(st, ed)的方向，单位为度。以东为正方向，逆时针为正（0 ~ 180），顺时针为负（0 ~ -180）
        """
        return np.arctan2(ed[1] - st[1], ed[0] - st[0]) / np.pi * 180

    def get_node_highest_class(self, road_pk, node_id):
        """
        返回node_id连接的路段中，除road_pk之外等级最高的道路级别。

        Args:
            road_pk: 路段主键
            node_id: 结点编号

        Returns:
            路段等级，如'C'。若node_id只连接road_pk则返回None
        """
        # 查看是否有依赖路段
        node_df = self.node_table.select_by_attr(self.node_table.pk, node_id)
        superior_class = node_df.loc[0, self.node_table.superior_class]

        # case 1:有依赖路段，返回依赖路段的最高等级
        if superior_class is not None:
            return superior_class

        # case 2：没有依赖路段，检查连接到node_id的其他路段。返回其他路段中的最高等级。
        prefix = self.road_table.name[:-2]
        scan_order = R_CLASS_LIST
        df = pd.DataFrame()
        for r_class in scan_order:
            df = pd.read_sql(f"""select * from {prefix}_{r_class.lower()}
                where {self.road_table.snodeid} = {node_id} or {self.road_table.enodeid}={node_id}""", conn)
            if r_class.lower() == self.road_table.name[-1].lower():
                df = df[df[self.road_table.pk] != road_pk]  # 去除自身
            if not df.empty:
                break
        if df.empty:  # node_id悬挂
            return None
        else:
            return r_class


def compute_stakes(input_shp_path, output_shp_path, road_prefix, node_table_name):
    gdf_pt = geopd.read_file(input_shp_path)
    for idx in gdf_pt.index:
        ldbm = gdf_pt.loc[idx, "ldbm"]
        r_class = ldbm[0]
        rt = FusedRoadTable(f"{road_prefix}_{r_class}")
        nt = RSNodeTable(node_table_name)
        st = StakeTool(rt, nt, 10)
        pt_geom = gdf_pt.loc[idx, "geometry"]
        stake_ = st.stake_interpolate(ldbm, pt_geom)
        print(stake_)
        gdf_pt.loc[idx, "xcoord"] = pt_geom.coords[0][0]
        gdf_pt.loc[idx, "ycoord"] = pt_geom.coords[0][1]
        gdf_pt.loc[idx, "stake"] = stake_
    gdf_pt.to_file(output_shp_path)


def label_stake_direction(road_prefix, node_table, east_west_angle):
    for r_class in R_CLASS_LIST:
        rt = FusedRoadTable(f"{road_prefix}_{r_class}".lower())
        st = StakeTool(rt, node_table, east_west_angle)
        gid_list = rt.get_pk_list()
        for gid in gid_list:
            print(r_class, gid)
            st.label_stake_direction(gid)


def get_all_stakes(pt_table_name, r_class_list, prefix):
    # 提取端点
    with conn.cursor() as cursor:
        cursor.execute(f"""
            drop table if exists {pt_table_name};
            create table {pt_table_name}(
            gid serial primary key,
            ldbm varchar(20),
            lxbm varchar(20),
            stake float,
            stake_rule varchar(20),
            stake_code varchar(2),
            geom geometry(Point,{SRID})
            )""")
        for r_class in r_class_list:
            table = FusedRoadTable(f"{prefix}_{r_class}".lower())
            gdf = geopd.read_postgis(f"select * from {table.name}", conn)
            for idx in gdf.index:
                geom = gdf.loc[idx, table.geom]
                st_pt = Point(geom.coords[0])
                ed_pt = Point(geom.coords[-1])
                ldbm = gdf.loc[idx, table.segment_code]
                lxbm = gdf.loc[idx, table.road_code]
                stake_rule = gdf.loc[idx, table.stake_rule]
                qdzh = gdf.loc[idx, table.start_stake]
                zdzh = gdf.loc[idx, table.end_stake]
                st_stake = qdzh
                ed_stake = zdzh
                cursor.execute(f"insert into {pt_table_name}(ldbm,lxbm,stake,stake_rule,stake_code,geom) "
                               f"values('{ldbm}','{lxbm}',{st_stake},'{stake_rule}','qd',{geom_to_ewkt(st_pt, SRID)})")
                cursor.execute(f"insert into {pt_table_name}(ldbm,lxbm,stake,stake_rule,stake_code,geom) "
                               f"values('{ldbm}','{lxbm}',{ed_stake},'{stake_rule}','zd', {geom_to_ewkt(ed_pt, SRID)})")
    conn.commit()
